#Author : Floriane Simonet
#e-mail : simonet.floriane@gmail.com
#! usr/bin/perl -w
use strict;


sub sort_by_chrpos{
	
	my $ref_chrpos=$_[0];
	my @chrpos=@$ref_chrpos;

	my @sortchrpos = sort {
	my @a = $a =~/(\d+)\-(\d+)/;
	my @b = $b =~/(\d+)\-(\d+)/; 
	$a[0] <=> $b[0] || $a[1] <=> $b[1] ;
	} @chrpos; # sort by chromosome then sort by position

	return(\@sortchrpos);
}

##########

sub create_snp_set2{

	my($ref_listchrp,$ref_genotypedsnpfile,$setfile1,$outputname)=@_;
	my @listchrp=@$ref_listchrp;
	my %snpfile=%$ref_genotypedsnpfile;

	my %groupset1=();
	my $currentgp=0;
	my $tmpos2;
	my $lastgp;

	open(SETFILE1,$setfile1) or die "cannot open SETFILE1 : $setfile1 $!\n";
	while(my $line=<SETFILE1>){
	
		chomp($line);
		my($chr,$pos,$rs,$gp)=split(/\t/,$line);

		if($currentgp==0){ # first group
			$currentgp=$gp;
			$lastgp=$gp;
			$groupset1{$gp}{'chr'}=$chr;
			$groupset1{$gp}{'pos1'}=$pos;
			$tmpos2=$pos;
		}
		elsif(($currentgp!=$gp) && ($currentgp!=0)){ #new group
			# add final information to group (n-1)
			$groupset1{$currentgp}{'pos2'}=$tmpos2;
			$groupset1{$currentgp}{'meanpos'}=sprintf("%0.f",((($groupset1{$currentgp}{'pos2'})+($groupset1{$currentgp}{'pos1'}))/2));
			$groupset1{$currentgp}{'limitdist'}=($groupset1{$currentgp}{'pos2'} - $groupset1{$currentgp}{'meanpos'}) * (2.5/1.5); # limits calculation
			if ( $groupset1{$currentgp}{'limitdist'} < 1000000 ){
				$groupset1{$currentgp}{'limitdist'}=1000000; # limit min
			}
			elsif ( $groupset1{$currentgp}{'limitdist'} > 2500000 ){
				$groupset1{$currentgp}{'limitdist'}=2500000; # limit max
			}
			$groupset1{$currentgp}{'limdown'}= $groupset1{$currentgp}{'meanpos'} - $groupset1{$currentgp}{'limitdist'};
			$groupset1{$currentgp}{'limup'}= $groupset1{$currentgp}{'meanpos'} + $groupset1{$currentgp}{'limitdist'};

			# create new group
			$currentgp=$gp;
			$lastgp=$gp;
			$groupset1{$gp}{'chr'}=$chr;
			$groupset1{$gp}{'pos1'}=$pos;
			$tmpos2=$pos;
		}
		else{ #same group
			$tmpos2=$pos;
		}
	}
	
	# add informations to last group	
	$groupset1{$lastgp}{'pos2'}=$tmpos2;
	$groupset1{$lastgp}{'meanpos'}=sprintf("%0.f",((($groupset1{$lastgp}{'pos2'})+($groupset1{$lastgp}{'pos1'}))/2));
	$groupset1{$lastgp}{'limitdist'}=($groupset1{$lastgp}{'pos2'} - $groupset1{$lastgp}{'meanpos'}) * (2.5/1.5); # limits calculation
			if ( $groupset1{$lastgp}{'limitdist'} < 1000000 ){
				$groupset1{$lastgp}{'limitdist'}=1000000; # limit min
			}
			elsif ( $groupset1{$lastgp}{'limitdist'} > 2500000 ){
				$groupset1{$lastgp}{'limitdist'}=2500000; # limit max
			}
	$groupset1{$lastgp}{'limdown'}= $groupset1{$lastgp}{'meanpos'} - $groupset1{$lastgp}{'limitdist'};
	$groupset1{$lastgp}{'limup'}= $groupset1{$lastgp}{'meanpos'} + $groupset1{$lastgp}{'limitdist'};
	
	close SETFILE1;

	my $cptup=0;
	my $cptdown=0;

	my @sortnum_gp= sort { $a <=> $b } keys%groupset1;

	#create directory for windowset files
	mkdir "$outputname/imputationSetfiles/windowsetfiles";

	# open warning file
	my $warnfile="$outputname/imputationSetfiles/warnfile_".$outputname.".txt";
	open(WARNFILE,">$warnfile") or die "cannot open WARNFILE : $warnfile $!\n";

	# open newbimfile to impute
	my $setfile2="$outputname/imputationSetfiles/genotyped_setfile_".$outputname.".txt";
	open(SETFILE2,">$setfile2") or die "cannot open SETFILE2 : $setfile2 $!\n";

	my $currentchr=0;
	my $numgpintochr;

	foreach my $num_gp (@sortnum_gp){

		my $firstline=1;
		my $pos1=0;
		my $pos2=0;		

		foreach my $k (@listchrp){ # read genotyped snpfile in chrpos order

			if(($snpfile{$k}{'chr'} == $groupset1{$num_gp}{'chr'}) && ($snpfile{$k}{'BPpos'} >= $groupset1{$num_gp}{'limdown'}) && ($snpfile{$k}{'BPpos'} <= $groupset1{$num_gp}{'limup'})){ # snps genotyped around target set (set of 7Mb max)

				# write genotyped_setfile
				print SETFILE2 "$snpfile{$k}{'chr'}\t$snpfile{$k}{'rsID'}\t$snpfile{$k}{'cMpos'}\t$snpfile{$k}{'BPpos'}\t$snpfile{$k}{'a1'}\t$snpfile{$k}{'a2'}\t$num_gp\n";

				# save information for windowsetfile
				$pos1 =$snpfile{$k}{'BPpos'} if($firstline == 1);
				$firstline=0;
				$pos2=$snpfile{$k}{'BPpos'};

				# count snps before target set
				if($snpfile{$k}{'BPpos'} < $groupset1{$num_gp}{'pos1'}){
					$cptdown=$cptdown+1;
				}
				
				# count snps after target set
				if($snpfile{$k}{'BPpos'} > $groupset1{$num_gp}{'pos2'}){
					$cptup=$cptup+1;
				}
			}
		}

		my $size = $pos2 - $pos1;

		if($currentchr!=$groupset1{$num_gp}{'chr'}){
			my $windowsetfile="$outputname/imputationSetfiles/windowsetfiles/windowset_".$outputname."_chr".$groupset1{$num_gp}{'chr'}.".txt";
			open(WINDOWFILE,">$windowsetfile") or die "cannot open WINDOWFILE : $windowsetfile $!\n";
			$currentchr=$groupset1{$num_gp}{'chr'};
			$numgpintochr=1;
		}

		print WINDOWFILE "$numgpintochr\t$pos1\t$pos2\t$size\n";
		$numgpintochr++;

		if($currentchr!=$groupset1{$num_gp}{'chr'}){
			close WINDOWFILE;
			$currentchr++;
		}

		# write warning file
		if(($cptup <= 10) || ($cptdown <= 10)){
			print WARNFILE " chr$groupset1{$num_gp}{'chr'} group$numgpintochr, not enough SNPs\n";
		} 
	}

	close SETFILE2;
	close WARNFILE;
	
}

##########

sub create_setfiles{

	my ($setfile1,$outputname)=@_;	
	my $currentgp=0;
	my $currentchr=0;
	my $numgpintochr=0;

	mkdir "$outputname/imputationSetfiles/rslist";

	open(SETFILE1, $setfile1) or die "cannot open $setfile1";
	while(my $line=<SETFILE1>)
	{
		chomp($line);
		my($chr,$BPpos,$rsID,$gp)=split(/\t/,$line); 

		$numgpintochr=1 if($currentchr!=$chr);		
		$currentchr=$chr;

		if($gp!=$currentgp){
			close SETFILE if($currentgp!=0); # not first line
			$currentgp++;

			my $setfile="$outputname/imputationSetfiles/rslist/target_setfile_".$outputname."_chr".$chr.".".$numgpintochr.".txt";
			open(SETFILE,">$setfile") or die "cannot open SETFILE : $setfile $!\n";
			print SETFILE "$rsID\n";
			$numgpintochr++;
		}
		else{
			print SETFILE "$rsID\n";
		}
	}	
	close SETFILE;
	close SETFILE1;
}

##########

sub create_snp_set1{

	my($ref_listchrp,$ref_snpfile,$outputname,$window_lg)=@_;

	my @listchrp=@$ref_listchrp;
	my %snpfile=%$ref_snpfile;

	my $gp=0;
	my $D=0;
	my $pos1=0;

	mkdir "$outputname/imputationSetfiles";

	my $chrlistfile="$outputname/imputationSetfiles/chr_list_".$outputname.".txt";
	open(CHRLIST,">$chrlistfile") or die "cannot open CHRLIST : $chrlistfile $!\n";

	my $setfile1="$outputname/imputationSetfiles/target_setfile_".$outputname.".txt";
	open(SETFILE1,">$setfile1") or die "cannot open SETFILE1 : $setfile1 $!\n";

	for(my $j=0; $j < scalar(@listchrp); $j++){ 

		my $chp=$listchrp[$j]; 

		if((exists($snpfile{$listchrp[$j-1]})) && ($snpfile{$chp}{'chr'} == $snpfile{$listchrp[$j-1]}{'chr'})){ # same chr
			$D=$pos1 - $snpfile{$chp}{'BPpos'};
			if( ((($snpfile{$listchrp[$j]}{'BPpos'}) - ($snpfile{$listchrp[$j-1]}{'BPpos'})) < 1000000) && ($D < $window_lg) ){ # si la distance du marqueur avec celui d'avant est < 1MB et si la distance du premier marqueur du groupe à ce marqueur est < $window_lg
				print SETFILE1 "$snpfile{$chp}{'chr'}\t$snpfile{$chp}{'BPpos'}\t$snpfile{$chp}{'rsID'}\t$gp\n";
			}
			else{

				$gp++;
				print SETFILE1 "$snpfile{$chp}{'chr'}\t$snpfile{$chp}{'BPpos'}\t$snpfile{$chp}{'rsID'}\t$gp\n";
				$pos1=$snpfile{$chp}{'BPpos'};
			}
		}
		elsif ((exists($snpfile{$listchrp[$j-1]})) && ($snpfile{$chp}{'chr'} != $snpfile{$listchrp[$j-1]}{'chr'} ) || (($j-1) == -1)){ # new chr

			$gp++;
			print SETFILE1 "$snpfile{$chp}{'chr'}\t$snpfile{$chp}{'BPpos'}\t$snpfile{$chp}{'rsID'}\t$gp\n";
			$D=$snpfile{$chp}{'BPpos'};
			print CHRLIST "$snpfile{$chp}{'chr'}\n";
			$pos1=$snpfile{$chp}{'BPpos'};
		}	
	}

	&create_setfiles($setfile1,$outputname);

	close(SETFILE1);
	close(CHRLIST);
	return ($setfile1);	
}

##########

sub read_SNPselectfile{

	my $SNPselectfile=$_[0];
	my $outputname=$_[1];
	my $window_lg=$_[2];
	$window_lg=$window_lg*1000000;

	my %snpfile=();
	my @chrpos=();

	open(SNPSELECTFILE, $SNPselectfile) or die "cannot open $SNPselectfile";
	while(my $line=<SNPSELECTFILE>)
	{
		chomp($line);
		my($chr,$rsID,$cMpos,$BPpos,$a1,$a2)=split(/\t/,$line); # read snp file
		
		my $chp=$chr."-".$BPpos; # get chr-position
		push(@chrpos,$chp);

		$snpfile{$chp}{'chr'}=$chr;
		$snpfile{$chp}{'rsID'}=$rsID;
		$snpfile{$chp}{'cMpos'}=$cMpos;
		$snpfile{$chp}{'BPpos'}=$BPpos;
		$snpfile{$chp}{'a1'}=$a1;
		$snpfile{$chp}{'a2'}=$a2;

	}
	close SNPSELECTFILE;

	my $ref_listchrp=&sort_by_chrpos(\@chrpos);
	my @listchrp=@$ref_listchrp;

	my ($setfile1)=&create_snp_set1(\@listchrp,\%snpfile,$outputname,$window_lg);

	return ($setfile1);
}

##########

sub read_SNPgenotypedfile{
	
	my $filename=$_[0];
	my $ref_snpset1=$_[1];
	my $outputname=$_[2];

	my %file=();
	my @chrpos=();

	open(FILE, $filename) or die "cannot open $filename";
	while(my $line=<FILE>)
	{
		chomp($line);
		my($chr,$rsID,$cMpos,$BPpos,$a1,$a2)=split(/\t/,$line); # read snp file
		
		my $chp=$chr."-".$BPpos; # get chr-position
		push(@chrpos,$chp);

		$file{$chp}{'chr'}=$chr;
		$file{$chp}{'rsID'}=$rsID;
		$file{$chp}{'cMpos'}=$cMpos;
		$file{$chp}{'BPpos'}=$BPpos;
		$file{$chp}{'a1'}=$a1;
		$file{$chp}{'a2'}=$a2;	
	}
	close FILE;

	my $ref_listchrp=&sort_by_chrpos(\@chrpos);
	my @listchrp=@$ref_listchrp;

	&create_snp_set2(\@listchrp,\%file,$ref_snpset1,$outputname);
}

############################################################

sub main(){
	
	if(@ARGV != 4){
		print "ERROR : missing parameters.\n Command line example :	perl create_imputation_setfiles.pl  target_snp_file  genotyped_snp_file  dataname length_target_snp_window\n";
		exit;
	}

	my $SNPselectfile=$ARGV[0]; # SNPs selected for replication (target SNP)
	my $SNPgenotypedfile=$ARGV[1];	
	my $outputname=$ARGV[2];
	my $window_lg=$ARGV[3]; 

	mkdir "$outputname";

	if(-f $SNPselectfile){
		my ($setfile1)=&read_SNPselectfile($SNPselectfile,$outputname,$window_lg);
		&read_SNPgenotypedfile($SNPgenotypedfile,$setfile1,$outputname);
	}
	else{
			print "$SNPselectfile doesn't exist.\n";
	}
exit;
}

main();

# command line example : perl create_imputation_setfiles.pl newbimfile_proj_SB.input_10252011.out newbimfile_SB.bim_1142011.out projSB 3
